#!/usr/bin/env bash
#
# Copyright (C) Hygon Info Technologies Ltd.
#
# SPDX-License-Identifier: Apache-2.0
#

set -o errexit
set -o errtrace
set -o pipefail

readonly script_name="$(basename "${BASH_SOURCE[0]}")"
readonly script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
readonly PKG_DIR="${script_dir}/../pkg"
mkdir -p "${PKG_DIR}"

readonly DTK_URL="https://download.sourcefind.cn:65024/directlink/1/DTK-24.04.1/Ubuntu20.04.1/DTK-24.04.1-Ubuntu20.04.1-x86_64.tar.gz"
readonly RCOK_URL="https://download.sourcefind.cn:65024/directlink/6/latest%E9%A9%B1%E5%8A%A8/previous_package/rock-5.7.1-6.2.18-V1.1.2.aio.run"

readonly hygon_vendor_id="0x1d94"

declare device
declare device_id
declare vendor_id

lib_file="${script_dir}/lib.sh"
source "$lib_file"


usage() {
    error="${1:-0}"
    cat <<EOF
Overview:

    Bound device to vfio for passthrough, and setup vfio device plugin for k8s.

Usage:

    $script_name [options]

Options:

    -h				: Display this help.
    -d <device> 	: Device need to be passthrough to guest
                      Examples: -b 0000:03:00.0
                      Hint: To find devices, use command:
                        lspci -vnn | grep Display

EOF
    exit "${error}"
}

parse_arguments() {
    [ "$#" -eq 0 ] && usage && return 0
    while getopts hd: opt
    do
        case $opt in
            d)	device=$OPTARG;;
            *)  die "Found an invalid option";;
        esac
    done

    if [ -z "${device}" ]; then
        die "Need device!"
    fi

    shift $(($OPTIND - 1))
}

bind_devices_to_vfio() {
    vendor_id=$(cat /sys/bus/pci/devices/$device/vendor)
    if [ ${vendor_id} != ${hygon_vendor_id} ]; then
        die "The device with vendor ${vendor_id} is not a hygon device!"
    fi

    device_id=$(cat /sys/bus/pci/devices/$device/device)
    if [ -e /sys/bus/pci/devices/$device/driver ]; then
        echo $device > /sys/bus/pci/devices/$device/driver/unbind
    fi

    echo $vendor_id $device_id > /sys/bus/pci/drivers/vfio-pci/new_id
    OK "device $device bound to vfio"
}

install_vfio_device_plugin() {
    export KUBECONFIG=/etc/kubernetes/admin.conf
    #Remove extra '0x'
    device_id=${device_id#"0x"}
    vendor_id=${vendor_id#"0x"}
    info "Adding $vendor_id $device_id to vfio device plugin"
    cat <<-EOF | kubectl apply -f -
apiVersion: v1
kind: ConfigMap
metadata:
  name: vfio-device-plugin-config
  namespace: kube-system
data:
  config.yml: |
    - resourceName: dcu/a100
      vendorId: ${vendor_id}
      deviceId: ['${device_id}']
---
apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: vfio-device-plugin
  namespace: kube-system
spec:
  selector:
    matchLabels:
      name: vfio-device-plugin
  template:
    metadata:
      annotations:
        scheduler.alpha.kubernetes.io/critical-pod: ""
      labels:
        name: vfio-device-plugin
    spec:
      tolerations:
      - key: CriticalAddonsOnly
        operator: Exists
      containers:
      - name: vfio-device-plugin
        image: docker.io/pawsonfang/vfio-device-plugin
        securityContext:
          allowPrivilegeEscalation: false
          capabilities:
            drop: ["ALL"]
        volumeMounts:
          - name: device-plugins
            mountPath: /var/lib/kubelet/device-plugins
          - name: dev
            mountPath: /dev
          - name: sys
            mountPath: /sys
          - name: config
            mountPath: /root/config
      volumes:
        - name: device-plugins
          hostPath:
            path: /var/lib/kubelet/device-plugins
        - name: dev
          hostPath:
            path: /dev
        - name: sys
          hostPath:
            path: /sys
        - name: config
          configMap:
            name: vfio-device-plugin-config
EOF
}

install_dcu_toolkit() {
    # DTK is too big to put it in ISO.
    # Download DTK-24.04.1 from https://cancon.hpccube.com:65024/1/main/DTK-24.04.1.
    # User should find the correct DTK release for your OS.
    # Here is a example for ubuntu 20.04.1
    [ ! -f DTK-24.04.1.tar.gz ] && wget -O DTK-24.04.1.tar.gz $DTK_URL
    tar xvf DTK-24.04.1.tar.gz -C /opt/
    echo 'export HSA_USERPTR_FOR_PAGED_MEM=0' >> /opt/dtk-24.04.1/env.sh
    echo 'hy-virtual -enable-encrypted' >> /opt/dtk-24.04.1/env.sh
}

# To make guest image as small as possible,
# install hyhal on host, then mount it to guest.
install_dcu_hyhal() {
    [ ! -f rock.aio.run ] && wget -O rock.aio.run $RCOK_URL
    chmod +x rock.aio.run
    ./rock.aio.run -u 0 -r 0

    # Disable hymgr on host so that DCU dirver will not be loaded.
    systemctl disable hymgr

    # replace hymgr with a new one
    [[ -d /opt/hyhal && ! -f /opt/hyhal/bin/hymgr-kata.bkp ]] && cp -v /opt/hyhal/bin/hymgr /opt/hyhal/bin/hymgr-kata.bkp

    if [ ! -f ${PKG_DIR}/hymgr ]; then
      /opt/hygon/tools/sftp download /hygon-devkit/4.0/2025-01-01/csv/confidential-containers/pkg/hymgr ${PKG_DIR}/hymgr
    fi
    
    cp -v ${PKG_DIR}/hymgr /opt/hyhal/bin/
}

main() {
    check_root
    parse_arguments $*

    bind_devices_to_vfio

    install_vfio_device_plugin
    install_dcu_toolkit
    install_dcu_hyhal

    OK "Environment is set up, enjoy!"
}

main $*
